# -*- coding: utf-8 -*-
"""
@Time       :   2024/06/13
@Author     :   Li Kuangyuan, Yi Junquan
@Version    :   1.0
@Contact    :   1767958859@qq.com
@Software   :   VsCode
"""

import csv
import os
import pickle

def read_data(dataset, quasi_index_list, category_flag_list):
    """从csv文件中读取数据

    Args:
        dataset (str): 当前准备读取的数据集名称
        quasi_index_list (list): 数据集中, 准标识属性的index集合
        category_flag_list (list): 是否为类别

    Returns:
        (list,list): (数据列表，表头)
    """    
    # 将获得文件所在的目录路径，以方便使用相对路径
    base_dir = os.path.dirname(os.path.abspath(__file__))

    statistics_dict = [{} for _ in quasi_index_list]  # 构造每列数据对应的统计字典(统计数值类型的数据)
    data_list = []  # 数据列表

    with open(os.path.join(base_dir, "../datasets/{}/{}.csv".format(dataset, dataset))) as file:
        # 以;为分割符读取文件
        reader = csv.reader(file, delimiter=";")

        for i, row_data in enumerate(reader):
            if i == 0:  # 读取第一行的表头信息
                header_list = row_data
            else:   # 数据部分的读取
                data_list.append(row_data)
                for i, index, flag in zip(range(len(quasi_index_list)), quasi_index_list, category_flag_list):
                    if not flag:
                        # 统计种类出现次数
                        statistics_dict[i][row_data[index]] =  statistics_dict[i].get(row_data[index], 0) + 1
        # 若目录不存在, 则创建目录
        if not os.path.exists(os.path.join(base_dir, "../results/numeric")):            
            os.makedirs(os.path.join(base_dir, "../results/numeric"))   
        # 按值对数值类型的列统计信息进行排序, 并输出成pickle文件进行记录
        for i, index, flag in zip(range(len(quasi_index_list)), quasi_index_list, category_flag_list):
            if not flag:
                with open(os.path.join(base_dir, "../results/numeric", dataset + '_' + header_list[index] + '_static.pickle'), 'wb') as file:
                    pickle.dump((statistics_dict[i], sorted(statistics_dict[i])), file)

    return (data_list, header_list)

def column_change(data_list, quasi_index_list):
    """对列进行重排序, 将准标识符属性放到前面, 其余属性放到后面

    Args:
        data_list (list): 数据列表
        quasi_index_list (list): 准标识符属性在原始数据中的序号

    Returns:
        list: 重新排序后的结果
    """    
    res = []
    for data in data_list:  # 逐行处理
        quasi_list = []
        not_quasi_list = []
        # 准标识符属性一个列表，非准标识符属性一个列表
        for i, elem in enumerate(data):
            if i in quasi_index_list:
                quasi_list.append(elem)
            else:
                not_quasi_list.append(elem)
        # 将准标识符属性放到前面, 其余属性放到后面
        quasi_list.extend(not_quasi_list)
        res.append(quasi_list)
    return res


def column_restore(data_list, quasi_index_list):
    """column_change的逆操作, 恢复列在原始数据中的排列

    Args:
        data_list (list): 数据列表
        quasi_index_list (list): 准标识符属性在原始数据中的序号

    Returns:
        list: 恢复以后的结果
    """    
    res = []
    for data in data_list:
        # 后半段，非准标识符属性
        not_quasi_list = data[len(quasi_index_list):]
        # 前半段，准标识符属性
        quasi_list = data[:len(quasi_index_list)]
        # 遍历准标识符属性, 将其插回原位
        for i, attr in enumerate(quasi_list):
            not_quasi_list.insert(quasi_index_list[i], attr)
        res.append(not_quasi_list)
    return res


def trans_to_col_priotity(data_list):
    """行列转换(类似行优先变列优先), 将data_list中的行和列互换, 同一列的数据使用一个list存储

    Args:
        data_list: 数据列表, 结构可以看作是一个二维数组

    Returns:
        list: 转换后的数据
    """    
    res = [[] for _ in data_list[0]]
    for data in data_list:
        for i, column in enumerate(data):
            # 将数据加入对应的列的list
            res[i].append(column)
    return res

def write_anonymized_data(anon_data, anonymity_method, header, k, dataset):
    """将匿名化的数据写入csv文件

    Args:
        anon_data (list): 匿名化的数据
        header (list): 数据表头
        k (int): k
        dataset (str): 被匿名化的数据集名称

    Returns:
        int: 匿名化的数据量
    """    
    # 将获得文件所在的目录路径，以方便使用相对路径
    base_dir = os.path.dirname(os.path.abspath(__file__))
    anon_dir = os.path.join(base_dir, "../results/anonymized_dataset", dataset, anonymity_method)
    if not os.path.exists(anon_dir):
        os.makedirs(anon_dir)  # 若目录不存在，则创建目录
    # 将数据转化成列表
    if isinstance(anon_data, dict):
        anon_data = anon_data.values()
    else:
        anon_data = [anon_data]
    # 写入数据，返回写入数量
    for i, data in enumerate(anon_data):
        with open(os.path.join(anon_dir, "anonymized_" + str(k) + '_' + str(i) + ".csv"), mode='w', newline='') as file:
            writer = csv.writer(file, delimiter=';')
            writer.writerow(header)
            writer.writerows(data)
    return len(anon_data)
